package org.wx.io.nio;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 多线程版本
 * 将读和写放入不同的线程中执行
 * @author Jaymie on 2021/2/18
 */
public class NioMultiThreadsV2 {
    private Selector selector;
    private final int port = 9090;

    private final ThreadPoolExecutor readPoolExecutor = new ThreadPoolExecutor(Runtime.getRuntime().availableProcessors(),
            Runtime.getRuntime().availableProcessors() + 2,
            60,
            TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(10),
            new ThreadFactoryBuilder().setNameFormat("read").build());

    private final ThreadPoolExecutor writePoolExecutor = new ThreadPoolExecutor(Runtime.getRuntime().availableProcessors(),
            Runtime.getRuntime().availableProcessors() + 2,
            60,
            TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(3000),
            new ThreadFactoryBuilder().setNameFormat("write").build());


    public static void main(String[] args) throws IOException {
        NioMultiThreadsV2 nioSingleThread = new NioMultiThreadsV2();
        nioSingleThread.start();
    }

    private void initServer() throws IOException {
        ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.bind(new InetSocketAddress(port));
        serverSocketChannel.configureBlocking(false);

        selector = Selector.open();
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);
    }

    private void start() throws IOException {
        initServer();
        System.out.println("start server on " + port);
        while (true) {
            System.out.println("waiting for ready event");
            int readyNum = selector.select();
            if (readyNum == 0) {
                continue;
            }

            Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
            if (iterator.hasNext()) {
                SelectionKey key = iterator.next();
                iterator.remove();

                if (key.isAcceptable()) {
                    ServerSocketChannel ssc = (ServerSocketChannel) key.channel();
                    SocketChannel socketChannel = ssc.accept();
                    ByteBuffer buffer = ByteBuffer.allocate(1024);
                    System.out.println("accept client connection: " + socketChannel.getRemoteAddress());
                    registerChannel(selector, socketChannel, SelectionKey.OP_READ, buffer);
                } else if (key.isReadable()) {
                    //key.cancel();
                    readPoolExecutor.execute(new ReadHandler(key, selector));
                } else if (key.isWritable()) {
                    //key.cancel();
                    writePoolExecutor.execute(new WriteHandler(key));
                }

            }
        }
    }

    public void registerChannel(Selector selector, SelectableChannel channel, int ops, ByteBuffer byteBuffer) throws IOException {
        if (channel == null) {
            return;
        }
        // 设置channel非阻塞
        channel.configureBlocking(false);
        channel.register(selector, ops, byteBuffer);
    }


    public static class WriteHandler implements Runnable {
        private final SelectionKey key;

        public WriteHandler(SelectionKey key) {
            this.key = key;
        }

        @Override
        public void run() {
            System.out.println(Thread.currentThread().getName() + " 开始执行写线程");
            try {
                SocketChannel socketChannel = (SocketChannel) key.channel();
                ByteBuffer byteBuffer = (ByteBuffer) key.attachment();
                // 往里面写数据
                byte[] bytes = "Hello world\n".getBytes(StandardCharsets.UTF_8);

                for (byte aByte : bytes) {
                    byteBuffer.put(aByte);
                    if (byteBuffer.position() == byteBuffer.limit()) {
                        byteBuffer.flip();
                        while (byteBuffer.hasRemaining()) {
                            socketChannel.write(byteBuffer);
                        }
                        byteBuffer.clear();
                    }
                }
                // 重新注册成为可读事件
                key.interestOps(SelectionKey.OP_READ);
                //socketChannel.register(key.selector(), SelectionKey.OP_READ, byteBuffer);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }


    /**
     * 读线程
     */
    public static class ReadHandler implements Runnable {
        private final SelectionKey key;
        private final Selector selector;

        public ReadHandler(SelectionKey key, Selector selector) {
            this.key = key;
            this.selector = selector;
        }

        @Override
        public void run() {
            System.out.println(Thread.currentThread().getName() + " 开始执行读线程");
            try {
                SocketChannel socketChannel = (SocketChannel) key.channel();
                ByteBuffer byteBuffer = (ByteBuffer) key.attachment();
                int count;
                byteBuffer.clear();

                // 从socketChannel 中往 byteBuffer 中写事件，读完了就是0
                while ((count = socketChannel.read(byteBuffer)) > 0) {
                    byteBuffer.flip();
                    System.out.println("content: " + StandardCharsets.UTF_8.decode(byteBuffer).toString());
                    byteBuffer.clear();
                }

                // 重新注册写事件
                System.out.println("开始注册");
                System.out.println(socketChannel.isOpen());
                key.interestOps(SelectionKey.OP_WRITE);
                //socketChannel.register(selector, SelectionKey.OP_WRITE, byteBuffer);
                System.out.println("注册为写事件");

                if (count < 0) {
                    // Close channel on EOF， invalidates the key
                    socketChannel.close();
                }
//                socketChannel.configureBlocking(false);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

}
